from src import const
import numpy as np
from src.utils import parse_args_and_merge_const
import torch
import torch.nn.functional as F
from src.dataset import Rescale, CenterCrop
from torchvision import transforms
import matplotlib.pyplot as plt
import skimage

class DrowsyDetector(object):

    def __init__(self, use_device='cpu'):
        print('load net...')
        net = const.USE_NET()
        net = net.to(const.device)
        net.load_state_dict(torch.load('models/' + const.MODEL_NAME, map_location={'cuda:0': use_device}))
        net.eval()
        self.net = net
        self.rescale = Rescale((168, 168))  # (128 + 208) / 2
        self.center_crop = CenterCrop(128)
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])
        print('ok..')

    def classify(self, cropped):
        cropped = self.rescale(cropped)
        cropped = self.center_crop(cropped)
        # 很奇怪的问题，在我自己的电脑上skimage rescale后会变成float类型
        # 用img_as_ubyte把它转回来
        cropped = skimage.img_as_ubyte(cropped)
        cropped = cropped.copy()
        self.raw_cropped = cropped

        cropped = self.to_tensor(cropped)
        cropped = self.normalize(cropped)
        cropped = torch.unsqueeze(cropped, 0)
        cropped = cropped.float()

        self.cropped = cropped
        output = self.net(cropped)

        att = None
        if hasattr(self.net, 'att'):
            att = torch.mean(self.net.att, dim=1)
            att = att[0, :].detach().cpu().numpy()
        _, predicted = torch.max(output.data, 1)
        prob = F.softmax(output, dim=1)
        # print('prob:')
        # print(prob.detach().cpu().numpy())
        predicted = predicted.detach().cpu().numpy()[0]
        prob = prob[0, predicted].item()
        return const.label_id2name[predicted], prob, att


if __name__ == '__main__':
    import librosa
    parse_args_and_merge_const()
    interface = WavRecgnition()
    samples, _ = librosa.load('data/raw/13307130239-01-02.wav', sr=const.SR)
    print(interface.classify(samples))
